import os
import torch
import numpy as np
import scipy.io
import scipy.io.wavfile
import scipy.signal
from sklearn.linear_model import LogisticRegression
import torch.nn as nn
from scipy.stats import kurtosis, skew
from patient_information import find_patient_files,load_patient_data,get_num_locations,get_locations,get_grade_number,cal_patient_acc
from My_Dataloader import Dataset2
from torch.utils.data import DataLoader, WeightedRandomSampler
from xgboost import XGBClassifier
from sklearn import svm
from sklearn.metrics import recall_score,accuracy_score, confusion_matrix, classification_report
from imblearn.under_sampling import RandomUnderSampler
from imblearn.over_sampling import SVMSMOTE,RandomOverSampler,SMOTE
import random

# init_seed=2023
# np.random.seed(init_seed) # 用于numpy的随机数
# # torch.backends.cudnn.benchmark = False
# random.seed(init_seed)
def cal_static(data):
    mean = np.mean(data, axis=0)
    feng = kurtosis(data, axis=0)
    # if np.isnan(feng).any:
    #     feng = np.zeros(3)
    pian = skew(data, axis=0)
    # if np.isnan(pian).any:
    #     pian = np.zeros(3)
    return mean,feng,pian

def cal_label(data):
    labels = np.argmax(data, axis=1)
    dist = np.zeros(3)
    for label in labels:
        dist[label] += 1
    return dist/len(labels)

#计算患者级别的统计特征和标签
def cal_vote_feat(path,all_id,all_y_pred,all_location):

    patient_files = find_patient_files(path)
    num_patient_files = len(patient_files)
    target = list()  # 存放病人的真是标签
    feat = list()  # 存放病人的预测标签
    num_correct_patient = 0
    for i in range(num_patient_files):
        current_patient_data = load_patient_data(patient_files[i])
        # num_location = get_num_locations(current_patient_data) #听诊区数目
        current_ID = current_patient_data.split(" ")[0]
        # current_locations = get_locations(current_patient_data)#[PV,TV,MV]
        current_patient_target = get_grade_number(current_patient_data)
        current_patient_AV = []
        current_patient_PV = []
        current_patient_TV = []
        current_patient_MV = []
        current_patient_Phc = []
        current_patient_mean = np.zeros((4, 3)) #输出概率均值
        current_patient_kurtosis = np.zeros((4, 3))#输出概率峰度
        current_patient_skewness = np.zeros((4, 3))#输出概率偏度
        current_patient_label = np.zeros((4, 3))#各类别占比
        # 计算当前ID下各个听诊区预测概率的和
        for n in range(len(all_id)):
            if all_id[n] == current_ID:
                if all_location[n] == 'AV':
                    current_patient_AV.append(all_y_pred[n])
                elif all_location[n] == 'PV':
                    current_patient_PV.append(all_y_pred[n])
                elif all_location[n] == 'TV':
                    current_patient_TV.append(all_y_pred[n])
                elif all_location[n] == 'MV':
                    current_patient_MV.append(all_y_pred[n])
                elif all_location[n] == 'Phc':
                    current_patient_Phc.append(all_y_pred[n])
        if len(current_patient_AV)>0:
            current_patient_AV = np.vstack(current_patient_AV)
            current_patient_mean[0],current_patient_kurtosis[0],current_patient_skewness[0]=cal_static(current_patient_AV)
            current_patient_label[0] = cal_label(current_patient_AV)
        if len(current_patient_PV) > 0:
            current_patient_PV = np.vstack(current_patient_PV)
            current_patient_mean[1],current_patient_kurtosis[1],current_patient_skewness[1]=cal_static(current_patient_PV)
            current_patient_label[1] = cal_label(current_patient_PV)
        if len(current_patient_TV) > 0:
            current_patient_TV = np.vstack(current_patient_TV)
            current_patient_mean[2],current_patient_kurtosis[2],current_patient_skewness[2]=cal_static(current_patient_TV)
            current_patient_label[2] = cal_label(current_patient_TV)
        if len(current_patient_MV) > 0:
            current_patient_MV = np.vstack(current_patient_MV)
            current_patient_mean[3],current_patient_kurtosis[3],current_patient_skewness[3]=cal_static(current_patient_MV)
            current_patient_label[3] = cal_label(current_patient_MV)
        # if len(current_patient_Phc) > 0:
        #     current_patient_Phc = np.vstack(current_patient_Phc)
        #     current_patient_mean[4],current_patient_kurtosis[4],current_patient_skewness[4]=cal_static(current_patient_Phc)
        #     current_patient_label[4] = cal_label(current_patient_Phc)
        current_patient_feat = np.hstack((current_patient_mean,current_patient_kurtosis,current_patient_skewness,current_patient_label))
        current_patient_feat = current_patient_feat.flatten()
        feat.append(current_patient_feat)
        target.append(current_patient_target)
    return np.vstack(feat),np.hstack(target)


def cal_patient(model, loader,id,location):
    val_acc = 0.0
    all_y_pred = []  # 存放3s样本的预测概率
    all_y_pred_label = []  # 存放3s样本的真实标签
    all_label = []  # 存放3s样本的预测标签
    all_id = []
    all_location = []

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = model.to(device)
    with torch.no_grad():
        model.eval()
        for i, data in enumerate(loader):
            soft= nn.Softmax(dim=1)
            x, y, z, x1, x2 = data
            x = x.to(device)
            y = y.to(device)
            z = z.to(device)
            x1 = x1.to(device)
            x2 = x2.to(device)
            input2 = torch.cat((x1, x2), dim=1)
            outputs = model(x, input2)
            outputs = soft(outputs)
            _, y_pred = outputs.max(1)
            all_y_pred.append(outputs.cpu().detach())
            all_label.append(y.cpu().detach())
            all_y_pred_label.append(y_pred.cpu().detach())
            # for ii in range(batch_size):
            # print(z)
            all_id.append(id[z.cpu().detach()])
            all_location.append(location[z.cpu().detach()])

    all_y_pred = np.vstack(all_y_pred)
    return all_id, all_y_pred, all_location

def replace_nan_with_zero(arr):
    arr[np.isnan(arr)] = 0
    return arr

def patient_feat_label(fold_path,model):
    data_path = os.path.join(fold_path, 'feature')
    label_path = os.path.join(fold_path, 'label')

    data_train = np.load(data_path + '/train_loggamma.npy', allow_pickle=True)
    data_vali = np.load(data_path + '/vali_loggamma.npy', allow_pickle=True)

    label_train = np.load(label_path + '/train_label.npy', allow_pickle=True)
    label_vali = np.load(label_path + '/vali_label.npy', allow_pickle=True)

    location_train = np.load(label_path + '/train_location.npy', allow_pickle=True)
    location_vali = np.load(label_path + '/vali_location.npy', allow_pickle=True)
    id_train = np.load(label_path + '/train_id.npy', allow_pickle=True)
    id_vali = np.load(label_path + '/vali_id.npy', allow_pickle=True)
    index_train = np.load(label_path + '/train_index.npy', allow_pickle=True)
    index_vali = np.load(label_path + '/vali_index.npy', allow_pickle=True)

    stat_feat_path = os.path.join(fold_path, 'statistical_feature')
    demo_train = np.load(stat_feat_path + '/train_demographic.npy', allow_pickle=True)
    demo_vali = np.load(stat_feat_path + '/vali_demographic.npy', allow_pickle=True)

    stat_train = np.load(stat_feat_path + '/train_static.npy', allow_pickle=True)
    stat_vali = np.load(stat_feat_path + '/vali_static.npy', allow_pickle=True)

    train_set = Dataset2(wav_label=label_train, wav_data=data_train, wav_index=index_train, demographic=demo_train,
                         static=stat_train)
    vali_set = Dataset2(wav_label=label_vali, wav_data=data_vali, wav_index=index_vali, demographic=demo_vali,
                        static=stat_vali)
    batch_size = 1
    train_loader = DataLoader(train_set, batch_size=batch_size)
    vali_loader = DataLoader(vali_set, batch_size=batch_size)

    trian_all_id, trian_all_y_pred, trian_all_location = cal_patient(model, train_loader, id_train, location_train)
    vali_all_id, vali_all_y_pred, vali_all_location = cal_patient(model, vali_loader, id_vali, location_vali)

    train_path = os.path.join(fold_path, "train_data")
    vali_path = os.path.join(fold_path, "vali_data")

    X_train, y_train = cal_vote_feat(train_path, trian_all_id, trian_all_y_pred, trian_all_location)
    X_test, y_test = cal_vote_feat(vali_path, vali_all_id, vali_all_y_pred, vali_all_location)
    patient_feature_path = os.path.join(fold_path, 'patient_feat')
    patient_label_path = os.path.join(fold_path, 'patient_label')
    if not os.path.exists(patient_feature_path):
        os.makedirs(patient_feature_path)
    if not os.path.exists(patient_label_path):
        os.makedirs(patient_label_path)

    np.save(patient_feature_path + r'/train_feat.npy', X_train)
    np.save(patient_feature_path + r'/test_feat.npy', X_test)
    np.save(patient_label_path + r'/train_label.npy', y_train)
    np.save(patient_label_path + r'/test_label.npy', y_test)

if __name__ == "__main__":
    fold_path = 'data/1_fold'
    # fold_path = 'data'
    model_result_path = os.path.join('ResNet_result',fold_path)
    model_dir = os.path.join(model_result_path, 'model/last_model')
    model = torch.load(model_dir)

    # patient_feat_label(fold_path, model)

    patient_feature_path = os.path.join(fold_path, 'patient_feat')
    patient_label_path = os.path.join(fold_path, 'patient_label')

    X_train = np.load(patient_feature_path + '/train_feat.npy', allow_pickle=True)
    X_test = np.load(patient_feature_path + '/test_feat.npy', allow_pickle=True)
    y_train = np.load(patient_label_path + '/train_label.npy', allow_pickle=True)
    y_test = np.load(patient_label_path + '/test_label.npy', allow_pickle=True)
    X_train = replace_nan_with_zero(X_train)
    X_test = replace_nan_with_zero(X_test)
    # 固定规则投票
    # patient_pre, patient_ture_label, patient_acc = cal_patient_acc(vali_path, vali_all_id, vali_all_y_pred, vali_all_location)

    sampler = RandomUnderSampler()
    # sampler = RandomOverSampler()
    # sampler = SVMSMOTE()
    # sampler = SMOTE()
    X_train, y_train = sampler.fit_resample(X_train, y_train)
    class_count = [0] * 3
    for label in y_train:
        class_count[label] += 1
    print("vali_set:", 'absent:', class_count[0], 'soft:', class_count[1], 'loud:', class_count[2])


    # model = svm.SVC(kernel='rbf', C=4)
    model = svm.SVC()
    # model = LogisticRegression()
    # model = XGBClassifier(n_estimators=40, random_state=2021, eval_metric=['logloss'])
    # model = XGBClassifier()
    #
    model.fit(X_train, y_train)
    y_pred = model.predict(X_test)
    # 计算平均准确率
    accuracy = accuracy_score(y_test, y_pred)
    print("Accuracy: {:.4f}".format(accuracy))

    # 计算混淆矩阵
    cm1 = confusion_matrix(y_test, y_pred)
    print("Confusion Matrix:\n", cm1)

    # 计算召回率、F1分数等
    # report = classification_report(y_test, y_pred)
    # print("Classification Report:\n", report)

    Absent_num = np.sum(cm1[0])
    Soft_num = np.sum(cm1[1])
    Loud_num = np.sum(cm1[2])
    Absent_recall_patient = cm1[0][0] / Absent_num
    Soft_recall_patient = cm1[1][1] / Soft_num
    Loud_recall_patient = cm1[2][2] / Loud_num
    Patient_UAR = (Absent_recall_patient + Soft_recall_patient + Loud_recall_patient) / 3
    print("------------------------------Patient result------------------------------")
    print("Absent_recall: %.4f, Soft_recall: %.4f, Loud_recall: %.4f,Patient_UAR: %.4f"
          % (Absent_recall_patient, Soft_recall_patient, Loud_recall_patient, Patient_UAR))
    a = np.sum(cm1, 0)
    Absent_Precision_patient = cm1[0][0] / a[0]
    Soft_Precision_patient = cm1[1][1] / a[1]
    Loud_Precision_patient = cm1[2][2] / a[2]

    Absent_f1_patient = (2 * Absent_recall_patient * Absent_Precision_patient) / (
            Absent_recall_patient + Absent_Precision_patient)
    Soft_f1_patient = (2 * Soft_recall_patient * Soft_Precision_patient) / (
            Soft_recall_patient + Soft_Precision_patient)
    Loud_f1_patient = (2 * Loud_recall_patient * Loud_Precision_patient) / (
            Loud_recall_patient + Loud_Precision_patient)
    Patient_f1 = (Absent_f1_patient + Soft_f1_patient + Loud_f1_patient) / 3
    print("Absent_F1: %.4f, Soft_F1: %.4f, Loud_F1: %.4f,Patient_F1: %.4f"
          % (Absent_f1_patient, Soft_f1_patient, Loud_f1_patient, Patient_f1))
